"""
模拟树形折叠的表格实现
使用 TableView 形式，但支持分组折叠/展开功能
"""

from PySide6.QtWidgets import (
    QApplication, QMainWindow, QTableView, QVBoxLayout, 
    QWidget, QPushButton, QHBoxLayout, QStyledItemDelegate, QStyle
)
from PySide6.QtGui import QStandardItemModel, QStandardItem, QColor, QFont, QPainter, QPixmap, QIcon
from PySide6.QtCore import Qt, QRect, QSize
import sys


class CollapseDelegate(QStyledItemDelegate):
    """折叠按钮代理 - 在第一列绘制折叠/展开图标"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        # 创建折叠/展开图标
        self.collapsed_icon = self._create_collapsed_icon()
        self.expanded_icon = self._create_expanded_icon()
    
    def _create_collapsed_icon(self):
        """创建折叠图标（右箭头 ▶）"""
        pixmap = QPixmap(16, 16)
        pixmap.fill(Qt.GlobalColor.transparent)
        painter = QPainter(pixmap)
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        
        # 绘制右箭头三角形
        from PySide6.QtGui import QPolygon, QBrush
        from PySide6.QtCore import QPoint
        triangle = QPolygon([
            QPoint(5, 4),   # 左上
            QPoint(5, 12),  # 左下
            QPoint(11, 8)   # 右中
        ])
        painter.setBrush(QBrush(QColor(100, 100, 100)))
        painter.setPen(Qt.PenStyle.NoPen)
        painter.drawPolygon(triangle)
        painter.end()
        return QIcon(pixmap)
    
    def _create_expanded_icon(self):
        """创建展开图标（下箭头 ▼）"""
        pixmap = QPixmap(16, 16)
        pixmap.fill(Qt.GlobalColor.transparent)
        painter = QPainter(pixmap)
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        
        # 绘制下箭头三角形
        from PySide6.QtGui import QPolygon, QBrush
        from PySide6.QtCore import QPoint
        triangle = QPolygon([
            QPoint(4, 5),   # 左上
            QPoint(12, 5),  # 右上
            QPoint(8, 11)   # 下中
        ])
        painter.setBrush(QBrush(QColor(100, 100, 100)))
        painter.setPen(Qt.PenStyle.NoPen)
        painter.drawPolygon(triangle)
        painter.end()
        return QIcon(pixmap)
    
    def paint(self, painter, option, index):
        """绘制单元格"""
        if index.column() != 0:
            super().paint(painter, option, index)
            return
        
        # 获取item数据
        item = index.model().itemFromIndex(index)
        if not item:
            super().paint(painter, option, index)
            return
        
        # 检查是否是分组行
        is_group = item.data(Qt.ItemDataRole.UserRole) == 'group'
        
        painter.save()
        
        # 绘制背景
        if option.state & QStyle.StateFlag.State_Selected:
            painter.fillRect(option.rect, option.palette.highlight())
        else:
            bg_color = index.data(Qt.ItemDataRole.BackgroundRole)
            if bg_color:
                painter.fillRect(option.rect, bg_color)
            else:
                painter.fillRect(option.rect, option.palette.base())
        
        if is_group:
            # 绘制折叠/展开图标
            is_expanded = item.data(Qt.ItemDataRole.UserRole + 1)
            icon = self.expanded_icon if is_expanded else self.collapsed_icon
            
            icon_rect = QRect(option.rect.x() + 5, option.rect.y() + (option.rect.height() - 16) // 2, 16, 16)
            icon.paint(painter, icon_rect)
            
            # 绘制文本（缩进）
            text_rect = option.rect.adjusted(25, 0, 0, 0)
            painter.setPen(option.palette.text().color())
            font = QFont()
            font.setBold(True)
            painter.setFont(font)
            painter.drawText(text_rect, Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter, item.text())
        else:
            # 普通数据行，缩进显示
            text_rect = option.rect.adjusted(35, 0, 0, 0)
            painter.setPen(option.palette.text().color())
            painter.drawText(text_rect, Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter, item.text())
        
        painter.restore()
    
    def sizeHint(self, option, index):
        """返回单元格大小"""
        size = super().sizeHint(option, index)
        return QSize(size.width(), max(size.height(), 28))


class CollapseTableView(QTableView):
    """支持折叠功能的表格视图"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.group_rows = {}  # 存储分组行信息 {group_name: {'row': row, 'expanded': True/False, 'children': [rows]}}
        self.data_row_counter = 0  # 数据行计数器（不包括分组行）
        
        # 隐藏表头
        self.horizontalHeader().hide()
        self.verticalHeader().hide()
        
        # 设置代理
        self.collapse_delegate = CollapseDelegate()
        self.setItemDelegateForColumn(1, self.collapse_delegate)  # 第1列（名称）
    
    def mousePressEvent(self, event):
        """处理点击事件 - 点击分组行时折叠/展开"""
        if event.button() == Qt.MouseButton.LeftButton:
            index = self.indexAt(event.pos())
            if index.isValid() and index.column() == 1:  # 第1列（名称）
                item = self.model().itemFromIndex(index)
                if item and item.data(Qt.ItemDataRole.UserRole) == 'group':
                    # 切换折叠状态
                    self.toggle_group(item)
                    event.accept()
                    return
        super().mousePressEvent(event)
    
    def toggle_group(self, group_item):
        """切换分组的折叠/展开状态"""
        group_name = group_item.text()
        
        if group_name not in self.group_rows:
            return
        
        group_info = self.group_rows[group_name]
        is_expanded = group_info['expanded']
        
        # 切换状态
        group_info['expanded'] = not is_expanded
        group_item.setData(not is_expanded, Qt.ItemDataRole.UserRole + 1)
        
        # 隐藏或显示子行
        for child_row in group_info['children']:
            self.setRowHidden(child_row, not is_expanded)
        
        # 刷新视图
        self.viewport().update()


class CollapseTableDemo(QMainWindow):
    """折叠表格演示主窗口"""
    
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.setup_model()
        self.setup_data()
    
    def init_ui(self):
        """初始化UI"""
        self.setWindowTitle('模拟折叠表格示例')
        self.setGeometry(100, 100, 800, 600)
        
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        
        # 按钮工具栏
        toolbar = QHBoxLayout()
        
        self.btn_expand_all = QPushButton('展开全部')
        self.btn_expand_all.clicked.connect(self.expand_all)
        toolbar.addWidget(self.btn_expand_all)
        
        self.btn_collapse_all = QPushButton('折叠全部')
        self.btn_collapse_all.clicked.connect(self.collapse_all)
        toolbar.addWidget(self.btn_collapse_all)
        
        self.btn_add_group = QPushButton('添加分组')
        self.btn_add_group.clicked.connect(self.add_group)
        toolbar.addWidget(self.btn_add_group)
        
        toolbar.addStretch()
        layout.addLayout(toolbar)
        
        # TableView
        self.table_view = CollapseTableView()
        self.table_view.setAlternatingRowColors(True)
        layout.addWidget(self.table_view)
    
    def setup_model(self):
        """设置数据模型"""
        self.model = QStandardItemModel(0, 6)  # 增加一列
        self.table_view.setModel(self.model)
        
        # 设置表头（隐藏但需要定义列）
        headers = ['No', '名称', '数量', '单价', '金额', '备注']  # 添加 No 列
        self.model.setHorizontalHeaderLabels(headers)
    
    def setup_data(self):
        """设置示例数据"""
        # 清空现有数据
        self.model.removeRows(0, self.model.rowCount())
        self.table_view.group_rows.clear()
        
        # 添加分组数据
        groups = [
            {
                'name': '电子产品',
                'children': [
                    ['笔记本电脑', '5', '5000', '25000', '畅销'],
                    ['平板电脑', '10', '3000', '30000', '热门'],
                    ['手机', '20', '4000', '80000', '新品'],
                ]
            },
            {
                'name': '家具',
                'children': [
                    ['办公桌', '15', '800', '12000', '促销'],
                    ['办公椅', '25', '500', '12500', ''],
                ]
            },
            {
                'name': '文具',
                'children': [
                    ['笔记本', '100', '10', '1000', ''],
                    ['钢笔', '50', '20', '1000', ''],
                    ['文件夹', '80', '5', '400', ''],
                ]
            }
        ]
        
        for group in groups:
            self.add_group_data(group['name'], group['children'])
        
        # 调整列宽
        self.table_view.resizeColumnsToContents()
        self.statusBar().showMessage('已加载3个分组，点击分组名称可折叠/展开')
    
    def add_group_data(self, group_name, children_data):
        """添加一个分组及其子数据"""
        # 添加分组行
        group_row = self.model.rowCount()
        
        # 第0列：空白（分组行不显示行号）
        no_item = QStandardItem('')
        no_item.setBackground(QColor(230, 240, 250))
        no_item.setFlags(Qt.ItemFlag.ItemIsEnabled)
        
        # 第1列：分组名称
        group_item = QStandardItem(group_name)
        group_item.setData('group', Qt.ItemDataRole.UserRole)
        group_item.setData(True, Qt.ItemDataRole.UserRole + 1)  # 默认展开
        group_item.setBackground(QColor(230, 240, 250))
        group_item.setFlags(Qt.ItemFlag.ItemIsEnabled)
        
        row_items = [no_item, group_item]
        for _ in range(4):  # 其他列
            empty_item = QStandardItem('')
            empty_item.setBackground(QColor(230, 240, 250))
            empty_item.setFlags(Qt.ItemFlag.ItemIsEnabled)
            row_items.append(empty_item)
        
        self.model.appendRow(row_items)
        
        # 添加子数据行
        child_rows = []
        for child_data in children_data:
            child_row = self.model.rowCount()
            child_rows.append(child_row)
            
            # 第0列：行号
            self.table_view.data_row_counter += 1
            no_item = QStandardItem(str(self.table_view.data_row_counter))
            no_item.setTextAlignment(Qt.AlignmentFlag.AlignCenter)
            
            row_items = [no_item]
            # 其他列：数据
            for value in child_data:
                item = QStandardItem(value)
                item.setTextAlignment(Qt.AlignmentFlag.AlignCenter)
                row_items.append(item)
            
            self.model.appendRow(row_items)
        
        # 记录分组信息
        self.table_view.group_rows[group_name] = {
            'row': group_row,
            'expanded': True,
            'children': child_rows
        }
    
    def expand_all(self):
        """展开所有分组"""
        for group_name, group_info in self.table_view.group_rows.items():
            if not group_info['expanded']:
                group_row = group_info['row']
                group_item = self.model.item(group_row, 0)
                self.table_view.toggle_group(group_item)
        
        self.statusBar().showMessage('已展开所有分组')
    
    def collapse_all(self):
        """折叠所有分组"""
        for group_name, group_info in self.table_view.group_rows.items():
            if group_info['expanded']:
                group_row = group_info['row']
                group_item = self.model.item(group_row, 0)
                self.table_view.toggle_group(group_item)
        
        self.statusBar().showMessage('已折叠所有分组')
    
    def add_group(self):
        """添加新分组"""
        import random
        group_name = f'新分组{random.randint(1, 100)}'
        children_data = [
            [f'商品{i}', str(random.randint(1, 50)), str(random.randint(10, 1000)), '', '']
            for i in range(3)
        ]
        
        self.add_group_data(group_name, children_data)
        self.table_view.resizeColumnsToContents()
        self.statusBar().showMessage(f'已添加分组：{group_name}')


def main():
    """主函数"""
    app = QApplication(sys.argv)
    window = CollapseTableDemo()
    window.show()
    sys.exit(app.exec())


if __name__ == '__main__':
    main()
